package cn.zk;

import org.I0Itec.zkclient.ExceptionUtil;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkNoNodeException;

import java.util.Collections;
import java.util.Comparator;
import java.util.List;

public class DistributedSimpleQueue<T> {

    protected final ZkClient zkClient;
    protected final String root;

    protected static final String NODE_NAME = "n_";

    public DistributedSimpleQueue(ZkClient zkClient,String root){
        this.zkClient = zkClient;
        this.root = root;
    }

    public int size(){
        return zkClient.getChildren(root).size();
    }

    public boolean isEmpty(){
        return zkClient.getChildren(root).size() == 0;
    }

    public boolean offer(T element){
        String nodeFullPath = root.concat("/").concat(NODE_NAME);
        try{
            zkClient.createPersistentSequential(nodeFullPath,element);
        }catch (ZkNoNodeException e) {
            zkClient.createPersistent(root);
            offer(element);
        }catch (Exception e){
            throw ExceptionUtil.convertToRuntimeException(e);
        }
        return true;
    }

    public T poll(){
        List<String> list = zkClient.getChildren(root);
        if(list.size() == 0){
            return null;
        }
        Collections.sort(list, new Comparator<String>() {
            @Override
            public int compare(String o1, String o2) {
                return getNodeNumber(o1,NODE_NAME).compareTo(getNodeNumber(o2,NODE_NAME));
            }
        });
        for(String nodeName : list){
            String nodeFullPath = root.concat("/").concat(nodeName);
            T node = zkClient.readData(nodeFullPath);
            zkClient.delete(nodeFullPath);
            return node;
        }
        return null;
    }

    private String getNodeNumber(String str,String nodeName){
        int index = str.lastIndexOf(nodeName);
        if(index >= 0){
            index += NODE_NAME.length();
            return index <= str.length() ? str.substring(index) : "";
        }
        return str;
    }
}
